import processFMCW as p
import matplotlib.pyplot as plt
import os

class Viewer:
    
    def __init__(self, folder=None):
        self.le = None
        self.ri = None
        self.fig1 = plt.figure(1, figsize=(12,8))
        self.ax1 = self.fig1.add_subplot(111)
        self.fig2 = plt.figure(2, figsize=(3,8))
        self.ax2 = self.fig2.add_subplot(111)
        self.cmin = 122
        self.cmax = 142  
        self.folder = folder
        self.procdir = 'F:\IPYdata\CBAND\\processed\\'
    
    
    def connect(self):
        self.cidgetleft = self.fig1.canvas.mpl_connect('key_press_event', self.getleft)
        self.cidgetright = self.fig1.canvas.mpl_connect('key_press_event', self.getright)
        self.cidremove = self.fig1.canvas.mpl_connect('key_press_event', self.remove)
        self.cidsave = self.fig1.canvas.mpl_connect('key_press_event', self.save)
        self.cidnext = self.fig1.canvas.mpl_connect('key_press_event', self.next)
        self.cidwigglewin = self.fig1.canvas.mpl_connect('motion_notify_event', self.wigglewin)
            
    
    def getList(self, filelist):
        self.file = iter(filelist)
        
    
    def loadData(self):
        file = self.file.next()
        if file[-4:-2] == '.g':
            FD = p.FGPR()
            FD.readRaw(file[:-4])
            self.TD = p.TGPR()
            self.TD.timeDomain(FD.Fdata)
            self.TD.staticCorrection()
            self.TD.wscs()
            if len(file[7:-4]) == 1:
                self.TD.filename = self.folder+'_0'+file[7:-4]+'.b'+file[-2:]
            else:
                self.TD.filename = self.folder+'_'+file[7:-4]+'.b'+file[-2:]
        else:
            self.TD = p.TGPR()
            self.TD.readBin(file, 805)
            self.TD.filename = file
        
        print self.TD.filename
        
        
    def draw(self):
        im1 = self.ax1.imshow(self.TD.Tdata, aspect='auto', vmin=self.cmin, vmax=self.cmax, cmap=plt.cm.binary, extent=[0, self.TD.traces, self.TD.t[-1], 0])
        self.ax1.set_xlabel('Trace #')
        self.ax1.set_ylabel('Time [ns]')
        self.ax1.set_title(self.TD.filename)
        self.fig1.canvas.draw()     
        
        
    def show(self):
        self.loadData()
        self.draw()
        plt.show()   
        
    
    ''' KEY EVENTS '''
        
    def getleft(self, event):
        if event.inaxes:
            if event.key == 'q':
                self.le = int(event.xdata)
                print self.le
            
            
    def getright(self, event):
        if event.inaxes:
            if event.key == 'w':
                self.ri = int(event.xdata)
                print self.ri
           

    def remove(self, event):
        if event.key == 'r':
            self.TD.removeTraces([self.le, self.ri])
            self.le = None
            self.ri = None
            self.draw()

    
    def save(self, event):
        if event.key == 's':
            self.TD.writeBin(self.procdir+self.TD.filename)
    
    
    def next(self, event):
        if event.key == 'n':
            self.ax1.clear()
            self.loadData()
            self.draw()
    
    
    ''' MOTION EVENTS '''
    
    def wigglewin(self, event):
        if event.inaxes:
            try:
                self.traceNo = int(event.xdata)
                plt.hold(False)
                plt.plot(self.TD.Tdata[:,self.traceNo],self.TD.t) 
                plt.hold(True)  
                plt.axhline(event.ydata) 
                plt.axis([self.cmin, self.cmax, self.ax1.axis()[2],  self.ax1.axis()[3]])
                plt.title('Trace # %i' % self.traceNo)  
                plt.xlabel('Power [dB]')
                plt.ylabel('Time [ns]')        
                self.fig2.canvas.draw()
            except OverflowError:
                print 'No valid data!'



def process():
    
    ''' First step is to get a list of all folders and files in the raw
    data directory '''
    import os
    rawdir = 'F:\\IPYdata\\CBAND'
    
    ext = '.gYY'
    folder = '071125' 
    
    proclist = []
    
    os.chdir(rawdir+'\\'+folder)
    filelist = os.listdir(os.getcwd())
    for file in filelist:
        if file[-4:] == ext:
            proclist.append(file)
    
    ''' Edit and store every single file '''
    VU = Viewer(folder)
    VU.connect()
    VU.getList(proclist)
    VU.show()


if __name__ == '__main__':
    pass
    
